import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

## x * x - 5 * y
def fx(x1, x2):
    r = np.sqrt(x1 ** 2 + x2 ** 2) + np.spacing(1)
    result = np.sin(r)/r
    return result


## x * x - 5 * y
def fitness(matrix):
    r = np.sqrt(matrix[:, 0] ** 2 + matrix[:, 1] ** 2) + np.spacing(1)
    result = np.sin(r) / r
    return result

# tournament selection
def selection(parent, offspring, k = 2):
    size, dim = parent.shape
    matrix = np.vstack((parent, offspring))
    reproduction = np.zeros((size, dim))
    for i in range(size):
        best_index = 0
        best_fitness = np.inf
        for j in range(k):
            index = np.random.randint(0, 2 * size)
            x = matrix[index,:].reshape((1, dim))
            fit = fitness(x)
            if best_fitness > fit:
                best_fitness = fit
                best_index = index
        reproduction[i,:] = matrix[best_index,:]
    return reproduction

# real-coded crossover
'''
parent: 父代
offspring: 子代
pc: 交叉概率
nc: 20<=nc<=30
'''
def crossover(parent, pc = 0.5, nc = 20):
    if np.random.random() < pc:
        beta = np.power((2 * pc), 1/(nc + 1))
    else:
        beta = np.power(1/(2 * (1 - pc)), 1/(nc + 1))
    size = parent.shape[0]
    reproduction = parent
    for i in range(0, size, 2):
        j = np.mod(i + 1, 2)
        reproduction[i,:] = 0.5 * ((1 + beta) * parent[i,:] +  (1 - beta) * parent[j,:])
        reproduction[j,:] = 0.5 * ((1 - beta) * parent[i,:] +  (1 + beta) * parent[j,:])
    return reproduction

# polynomial mutation
'''
parent: 父代
nm: 20<=nm<=30
xrange: 决策变量的范围
'''
def mutation(parent, xrange, nm = 20):
    mu = np.random.random()
    size, dim = parent.shape
    low = bl = xrange[:,0]
    high = bh = xrange[:,1]
    for i in range(1, size):
        low = np.vstack((low, bl))
        high = np.vstack((high, bh))

    delta1 = (parent - low) / (high - low)
    delta2 = (high - parent) / (high - low)

    if mu < 0.5:
        delta = np.power((2 * mu + (1 - 2 * mu) * np.power((1 - delta1), nm)),
                         1 / (nm + 1)) - 1
    else:
        delta = 1 - np.power((2 * (1 - mu) + 2 * (mu - 0.5) * np.power((1 - delta2), nm + 1)), 1 / (nm + 1))
    reproduction = parent + delta * (high - low)
    return reproduction

if __name__ == '__main__':
    size = 50
    iteration = 200
    # 初始化
    xrange = np.array([[-8, 8], [-8, 8]])
    # xrange = np.array([[7, 26]])
    dim = xrange.shape[0]
    parent = np.zeros((size, 1))
    for d in range(dim):
        xd = xrange[d][0] +  (xrange[d][1] - xrange[d][0]) * np.random.random((size, 1))
        parent = np.hstack((parent, xd))
    offspring = parent = np.delete(parent, 0, axis=1)
    fig = plt.figure()
    for i in range(iteration):
        print(i)
        produce = selection(parent, offspring)
        produce = crossover(produce)
        produce = mutation(produce, xrange)
        for j in range(size):
            if fitness(produce[j,:][np.newaxis,:]) > fitness(parent[j,:][np.newaxis,:]):
                produce[j,:] = parent[j,:]
        parent = offspring
        offspring = produce

        x1, x2 =  np.ogrid[xrange[0][0]:xrange[0][1]:200j,xrange[1][0]:xrange[1][1]:200j]
        extent = [np.min(x1), np.max(x1), np.min(x2), np.max(x2)]
        f = fx(x1, x2)

        ax = fig.add_subplot(projection='3d')
        ax.plot_surface(x1,x2,f,cmap=cm.ocean)
        ax.scatter(offspring[:,0], offspring[:,1], fx(offspring[:,0], offspring[:,1]), c = "red")
        ax.set_xlabel("x1")
        ax.set_ylabel("x2")
        ax.set_zlabel("fx")
        plt.draw()
        plt.pause(0.2)
        plt.clf()

    ax = fig.add_subplot(projection='3d')
    ax.plot_surface(x1, x2, f, cmap=cm.ocean)
    ax.scatter(offspring[:, 0], offspring[:, 1], fx(offspring[:, 0], offspring[:, 1]), c="red")
    print(np.min(fx(offspring[:, 0], offspring[:, 1])))
    plt.show()